// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/init/init.h>
#include <gtest/gtest.h>

#include <pollux/common/memory/memory.h>
#include <pollux/dwio/parquet/register_parquet_reader.h>
#include <pollux/dwio/parquet/register_parquet_writer.h>
#include <pollux/exec/table_writer.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/hive_connector_test_base.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo::pollux::plan;

namespace kumo::pollux::test {

class InsertTest : public pollux::VectorBuilder {
 public:
  void runInsertTest(
      std::string_view outputDirectory,
      int numRows,
      memory::MemoryPool* pool) {
    auto rowType = ROW(
        {"c0", "c1", "c2", "c3"}, {BIGINT(), INTEGER(), SMALLINT(), DOUBLE()});

    auto input = make_row_vector(
        {make_flat_vector<int64_t>(numRows, [](auto row) { return row; }),
         make_flat_vector<int32_t>(numRows, [](auto row) { return row; }),
         make_flat_vector<int16_t>(numRows, [](auto row) { return row; }),
         make_flat_vector<double>(numRows, [](auto row) { return row; })});

    // Insert with one writer.
    auto plan =
        PlanBuilder()
            .values({input})
            .tableWrite(
                outputDirectory.data(), dwio::common::FileFormat::PARQUET)
            .planNode();

    // Execute the write plan.
    auto results = exec::test::AssertQueryBuilder(plan).copyResults(pool);

    // First column has number of rows written in the first row and nulls in
    // other rows.
    auto rowCount = results->childAt(exec::TableWriteTraits::kRowCountChannel)
                        ->as<FlatVector<int64_t>>();
    ASSERT_FALSE(rowCount->is_null_at(0));
    ASSERT_EQ(numRows, rowCount->value_at(0));
    ASSERT_TRUE(rowCount->is_null_at(1));

    // Second column contains details about written files.
    auto details = results->childAt(exec::TableWriteTraits::kFragmentChannel)
                       ->as<FlatVector<StringView>>();
    ASSERT_TRUE(details->is_null_at(0));
    ASSERT_FALSE(details->is_null_at(1));
    melon::Dynamic obj = melon::parseJson(details->value_at(1));

    ASSERT_EQ(numRows, obj["rowCount"].asInt());
    auto fileWriteInfos = obj["fileWriteInfos"];
    ASSERT_EQ(1, fileWriteInfos.size());

    auto writeFileName = fileWriteInfos[0]["writeFileName"].asString();

    // Read from 'writeFileName' and verify the data matches the original.
    plan = PlanBuilder().tableScan(rowType).planNode();

    auto filePath = fmt::format("{}{}", outputDirectory, writeFileName);
    const int64_t fileSize = fileWriteInfos[0]["fileSize"].asInt();
    auto split = exec::test::HiveConnectorSplitBuilder(filePath)
                     .fileFormat(dwio::common::FileFormat::PARQUET)
                     .length(fileSize)
                     .build();
    auto copy =
        exec::test::AssertQueryBuilder(plan).split(split).copyResults(pool);
    exec::test::assertEqualResults({input}, {copy});
  }
};
} // namespace kumo::pollux::test
